#include <tinx/part.h>
#include <tinx/memory.h>
#include <tinx/device.h>
#include <tinx/debug.h>
#include <tinx/user.h>
#include <tinx/task.h>
#include <tinx/bitmap.h>
#include <tinx/stdlib.h>
#include <tinx/string.h>

#define LOGK(str, args...) DEBUGK(str, ##args)

#define PAGE_SECTORS (PAGE_SIZE / SECTOR_SIZE)

typedef struct swap_t
{
    list_t waiters;
    task_t *task;
    bitmap_t map;
    u32 used_pages;
    u32 pages; // 页数
    dev_t dev;
} swap_t;

static swap_t swapdev;

extern u8 *memory_map;
extern task_t *task_table[TASK_NR];

idx_t get_swap_page()
{
    swap_t *swap = &swapdev;
    idx_t idx = bitmap_scan(&swap->map, 1);
    swap->used_pages++;
    return idx;
}

void put_swap_page(idx_t idx)
{
    swap_t *swap = &swapdev;

    swap->used_pages--;

    assert(bitmap_test(&swap->map, idx));
    bitmap_set(&swap->map, idx, false);
}

int swap_read(idx_t idx, u32 vaddr)
{
    swap_t *swap = &swapdev;
    int ret = device_request(swap->dev, REQ_READ, (void *)vaddr, PAGE_SECTORS, idx * PAGE_SECTORS, 0);
    return ret;
}

int swap_write(idx_t idx, u32 vaddr)
{
    swap_t *swap = &swapdev;
    int ret = device_request(swap->dev, REQ_WRITE, (void *)vaddr, PAGE_SECTORS, idx * PAGE_SECTORS, 0);
    return ret;
}

void swap_in(swap_entry_t *entry)
{
    swap_t *swap = &swapdev;

    assert(swap->dev != EOF);
    assert(!entry->present);

    idx_t idx = entry->index;
    assert(idx);

    u32 page = get_page();
    entry_init((page_entry_t *)entry, page);

    u32 vaddr = pte_vaddr((page_entry_t *)entry);
    flush_tlb(vaddr);

    int ret = swap_read(idx, vaddr);
    assert(ret >= 0);
    put_swap_page(idx);
}

// 尝试把内存交换到 swap 设备中
// 成功返回 true 失败返回 false
bool try_to_swap_out(page_entry_t *entry)
{
    // 保证有 swap
    swap_t *swap = &swapdev;
    if (swap->dev == EOF)
        return false;
    assert(entry->present);
    u32 page = PAGE(entry->index);
    // 不能是内核页
    if (page < KERNEL_MEMORY_SIZE)
        return false;

    u32 vaddr = pte_vaddr(entry);
    if (entry->dirty)
    {
        // 不能是共享页
        if (memory_map[entry->index] != 1)
            return false;
        // 开始交换
        idx_t idx = get_swap_page();
        if (idx == EOF)
            return false;
        int ret = swap_write(idx, vaddr);
        assert(ret >= 0);
        swap_entry_t *sentry = (swap_entry_t *)entry;
        sentry->present = false;
        sentry->index = idx;
        flush_tlb(vaddr);
        // 释放原来的页
        assert(ret >= 0);
        put_page(page);
        return true;
    }
    // 直接释放
    entry_free(entry);
    flush_tlb(vaddr);
    put_page(page);
    return true;
}

// 尝试交换任务的所有页，返回成功交换掉的页数
// 为了简单起, 页表没有交换, 只交换页框
u32 task_swap_out(task_t *task)
{
    u32 old_cr3 = get_cr3();

    // 为了方便先切换页目录
    page_entry_t *pgdir = (page_entry_t *)task->pgdir;
    set_cr3((u32)task->pgdir);

    u32 start = DIDX(USER_EXEC_ADDR);
    u32 end = DIDX(USER_MEMORY_END);

    u32 swap_pages = 0;

    for (int didx = start; didx < end; didx++)
    {
        page_entry_t *pde = &pgdir[didx];

        // 不存在就不交换
        if (!pde->present)
            continue;

        // 获得页表
        page_entry_t *table = (page_entry_t *)(PDE_MASK | (didx << 12));

        for (int tidx = 0; tidx < 1024; tidx++)
        {
            page_entry_t *pte = &table[tidx];

            // 不存在就不交换
            if (!pte->present)
                continue;

            // 如果存在，尝试交换
            if (try_to_swap_out(pte))
                swap_pages++;
        }
    }

    // 恢复原来的页目录
    set_cr3(old_cr3);
    return swap_pages;
}

u32 swap_out()
{
    swap_t *swap = &swapdev;

    // 交换掉的页的数量
    u32 swap_pages = 0;
    if (swap->dev == EOF)
        return swap_pages;

    // 只允许一个任务执行此函数，其他任务等着领就行
    list_t *list = &swap->waiters;
    task_t *task = running_task();
    if (swap->task)
    {
        list_append(list, &task->node);
        task_block(task, TASK_BLOCKED);
        return 1;
    }
    swap->task = task;

    // 尝试交换所有任务的页
    for (int i = 0; i < TASK_NR; i++)
    {
        task_t *ptr = task_table[i];
        if (!ptr)
            continue;

        // 不允许交换内核任务的页
        if (ptr->uid == KERNEL_USER)
            continue;

        swap_pages += task_swap_out(ptr);
    }

    // 把交换掉的页分给等待的任务
    for (int i = 0; i < MIN((swap_pages - 1), list_size(list)); i++)
    {
        list_node_t *node = list_pop(list);
        assert(node);
        task_t *waiter = element_entry(task_t, node, node);
        task_unblock(waiter, EOK);
        swap_pages--;
    }
    swap->task = NULL;

    return swap_pages;
}

void swap_init()
{
    swap_t *swap = &swapdev;
    memset(swap, 0, sizeof(swap_t));
    swap->dev = EOF;

    for (int i = 1; true; i++)
    {
        dev_t dev = device_find(DEV_FLAG_PART, i);
        if (dev == EOF)
        {
            break;
        }
        int type = device_ioctl(dev, DEV_CMD_PART_TYPE, 0);
        if (type == PART_FS_SWAP)
        {
            swap->dev = dev;
            break;
        }
    }

    if (swap->dev == EOF)
    {
        return;
    }

    list_init(&swap->waiters);
    swap->task = NULL;

    swap->pages = device_ioctl(swap->dev, DEV_CMD_SECTORS_COUNT, 0) / PAGE_SECTORS;
    swap->used_pages = 0;

    LOGK("swap device number %d pages %d\n", swap->dev, swap->pages);

    bitmap_init(
        &swap->map,
        (u8 *)alloc_kpage(div_round_up(swap->pages / 8, PAGE_SIZE)),
        swap->pages / 8,
        1);
}